import pandas as pd
import numpy as np


def combine_data(file_path, sheet_names):
    dfs = []
    for i in range(len(sheet_names)):
        df = pd.read_excel(file_path, sheet_name=sheet_names[i])
        df['磁芯材料'] = sheet_names[i]
        dfs.append(df)
    # 合并为一个DataFrame
    combined_df = pd.concat(dfs, ignore_index=True)

    return combined_df


def data_process(df):
    # 波形
    wave_dict = {'正弦波': 1, '三角波': 2, '梯形波': 3}
    # 材料
    material_dict = {'材料1': 1, '材料2': 2, '材料3': 3, '材料4': 4}
    # 使用map()方法将列中的值映射到新值
    df['励磁波形'] = df['励磁波形'].map(wave_dict)
    df['磁芯材料'] = df['磁芯材料'].map(material_dict)


# 提取磁通密度的峰值、峰峰值、平均变化率
def flux_extract(df):
    # 初始化一个空的列表来存储特征
    features = []
    # 计算特征
    for index, row in df.iterrows():
        signal = row.values
        max_value = np.max(signal)  # 磁通密度最大值
        min_value = np.min(signal)  # 磁通密度最小值
        peak_to_peak = max_value - min_value  # 磁通密度峰峰值
        # 平均磁通密度变化率
        diff = np.diff(signal)
        diff_mean = np.mean(diff)
        # 将特征添加到列表中
        features.append({
            '磁通密度最大值': max_value,
            '磁通密度最小值': min_value,
            '磁通密度峰峰值': peak_to_peak,
            '平均磁通密度变化率': diff_mean
        })
    # 将特征列表转换为DataFrame
    features_df = pd.DataFrame(features)

    return features_df


if __name__ == '__main__':
    # 提取训练数据
    print("开始提取训练数据")
    train_file_path = '../2024年中国研究生数学建模竞赛赛题/C题/附件一（训练集）.xlsx'
    sheet_names = ['材料1', '材料2', '材料3', '材料4']
    combined_df = combine_data(train_file_path, sheet_names)
    data_process(combined_df)
    env_feature = combined_df.loc[:, ['温度，oC', '频率，Hz', '磁芯材料', '励磁波形']]
    flux_density = combined_df.iloc[:, 4:-1]
    core_loss = combined_df['磁芯损耗，w/m3']
    flux_feature = flux_extract(flux_density)
    train_data = pd.concat([env_feature, flux_feature, core_loss], axis=1)
    train_data.to_excel('../data/Q4_train.xlsx', index=False)
    print("训练数据提取完成")

    # 提取预测数据
    print("开始提取测试数据")
    test_file_path = '../2024年中国研究生数学建模竞赛赛题/C题/附件三（测试集）.xlsx'
    test_df = pd.read_excel(test_file_path)
    data_process(test_df)
    env_feature = test_df.loc[:, ['温度，oC', '频率，Hz', '磁芯材料', '励磁波形']]
    flux_density = test_df.iloc[:, 5:]
    flux_feature = flux_extract(flux_density)
    test_data = pd.concat([env_feature, flux_feature], axis=1)
    test_data.to_excel('../data/Q4_test.xlsx', index=False)
    print("测试数据提取完成")
